package com.zf.nospring.kafka.transaction;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;

import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Properties;
import java.util.concurrent.Future;

/**
 * 使用事务完成：消费->处理->生产。
 * 将 消费、生产 放在一个事务中，如果其中出现异常，那么就消费和生产的操作全部失败回滚。
 * <p>
 * 从 first 主题中读取数据，处理后，写入 second 主题中。整个操作在一个事务中，如果出现异常，则全部失败回滚。
 */
public class ConsumerProcessProducerTest {


    /**
     * 完成 "从 first 主题中读取数据，处理后，写入 second 主题中" 操作
     */
    public static void main(String[] args) {
        KafkaProducer<String, String> producer = createProducer();
        KafkaConsumer<String, String> consumer = createConsumer();

        // 初始化事务
        producer.initTransactions();
        try (producer; consumer) {
            while (true) {
                // 开启事务
                producer.beginTransaction();

                // 存储topic、partition和 偏移量offset
                HashMap<TopicPartition, OffsetAndMetadata> offsetMap = new HashMap<>();

                // 拉取一批消息
                ConsumerRecords<String, String> consumerRecords = consumer.poll(Duration.ofSeconds(1));
                // 处理消息，并发送到其他topic
                for (ConsumerRecord<String, String> record : consumerRecords) {
                    String topic = record.topic();
                    int partition = record.partition();
                    long offset = record.offset();
                    String content = record.value();

                    // 保存topic、partition 和 offset，用于手动提交；offset要加1 即指定下一条消息的偏移
                    offsetMap.put(new TopicPartition(topic, partition), new OffsetAndMetadata(offset + 1));

                    // 处理消息内容，添加 叹号
                    content = content + "!";

                    // 模拟异常
                    // int a = 1/0;

                    // 发送消息到 second 主题
                    Future<RecordMetadata> future = producer.send(new ProducerRecord<>("second", content));
                    try {
                        future.get();
                    } catch (Exception e) {
                        e.printStackTrace();
                        // 发生异常，终止事务，消费和生产的操作会回滚
                        producer.abortTransaction();
                    }
                }

                // 把偏移发送到当前事务中，发送的偏移量应该是将要消费的下一条消息即 offset+1，指定消费者组ID。
                // 会将这些偏移标记为当前事务的一部分，仅当事务成功提交时，这些偏移才会被视为已提交，也就是消费完成。否则就是消费失败，消息依然存在topic中
                producer.sendOffsetsToTransaction(offsetMap, new ConsumerGroupMetadata("dgc"));
                // 提交事务，确定发送完成、同时确定消费完成
                producer.commitTransaction();

            }
        } catch (Exception e) {
            e.printStackTrace();
            // 发生异常，终止事务，消费和生产的操作会回滚
            producer.abortTransaction();
        }

    }

    /**
     * 创建生产者，开启手动提交offset
     */
    private static KafkaConsumer<String, String> createConsumer() {
        // 1. 配置消费者的属性（添加对事务的支持）
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "dgc");
        // 配置事务的隔离级别
        props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
        // 关闭自动提交，一会我们需要手动来提交offset，通过事务来维护offset
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
        // 反序列化器
        props.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        props.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());

        // 2. 构建消费者对象
        KafkaConsumer<String, String> kafkaConsumer = new KafkaConsumer<>(props);

        // 3. 订阅一个topic
        kafkaConsumer.subscribe(Arrays.asList("first"));

        return kafkaConsumer;
    }

    /**
     * 创建一个带有事务配置的生产者
     */
    private static KafkaProducer<String, String> createProducer() {
        // 1. 配置生产者带有事务配置的属性
        Properties props = new Properties();
        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());

        // 开启事务必须要配置事务的ID
        props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "cpp_trans");
        // 2. 构建生产者
        return new KafkaProducer<>(props);
    }
}
